package logic

import (
	"context"
	"net"
	"time"

	"probe-tools/internal/svc"
	"probe-tools/pb/probe"

	"github.com/zeromicro/go-zero/core/logx"
)

type DnsProbeLogic struct {
	ctx    context.Context
	svcCtx *svc.ServiceContext
	logx.Logger
}

func NewDnsProbeLogic(ctx context.Context, svcCtx *svc.ServiceContext) *DnsProbeLogic {
	return &DnsProbeLogic{
		ctx:    ctx,
		svcCtx: svcCtx,
		Logger: logx.WithContext(ctx),
	}
}

func (l *DnsProbeLogic) DnsProbe(in *probe.DnsProbeRequest) (*probe.DnsProbeResponse, error) {
	start := time.Now()
	timeout := time.Duration(in.Timeout) * time.Second
	if in.Timeout <= 0 {
		timeout = 5 * time.Second // 默认超时5秒
	}

	// 设置超时上下文
	ctx, cancel := context.WithTimeout(l.ctx, timeout)
	defer cancel()

	// 解析DNS
	ips, err := net.DefaultResolver.LookupIP(ctx, "ip", in.Domain)
	if err != nil {
		return &probe.DnsProbeResponse{
			Success:  false,
			Duration: time.Since(start).Milliseconds(),
			Error:    err.Error(),
		}, nil
	}

	// 转换IP为字符串
	ipStrs := make([]string, 0, len(ips))
	for _, ip := range ips {
		ipStrs = append(ipStrs, ip.String())
	}

	return &probe.DnsProbeResponse{
		Success:  true,
		Ips:      ipStrs,
		Duration: time.Since(start).Milliseconds(),
		Type:     "dns",
	}, nil
}
